import numpy as np
from auxmodels.utils.calc import xyxy2xywh


def merge_cross_scales(results):
    """
    results为多个尺度视图的结果列表
    """
    results = [x for x in results if x is not None]
    if len(results) > 0:
        results = np.concatenate(results, axis=0)
        return results
    else:
        return


def merge_overlaps(results):
    """合并重叠框"""
    if results is not None and len(results) > 1:
        dist_ratios = []
        for obj in results:
            dist_ratios.append(calc_box_dist_ratio(obj[3:], results[:, 3:]))
        dist_ratios = np.array(dist_ratios, dtype=np.float32)
        for idx in range(len(dist_ratios)):
            dist_ratios[idx, idx] = 999  # 对角线设置一个大的值，屏蔽自身
        return cluster_boxes(results, dist_ratios)
    else:
        return results


def calc_box_dist_ratio(base, boxes):
    """计算目标框之间的距离"""
    base = xyxy2xywh(base)
    boxes = xyxy2xywh(boxes)
    diagonal_base = np.sqrt(base[2] ** 2 + base[3] ** 2)
    diagonal_boxes = np.sqrt(boxes[:, 2] ** 2 + boxes[:, 3] ** 2)

    cx0, cy0 = base[0], base[1]
    cx1, cy1 = boxes[:, 0], boxes[:, 1]

    dists = np.sqrt((cx1 - cx0) ** 2 + (cy1 - cy0) ** 2)
    ratios = dists / (diagonal_base / 2 + diagonal_boxes / 2)
    return ratios


def cluster_boxes(results, dist_ratios, thres=0.9):
    """根据框之间的距离矩阵，进行聚类"""
    clustered_list = []
    indexes = np.where(dist_ratios < thres)
    for index in zip(indexes[0], indexes[1]):
        if index[0] < index[1]:
            if (idx := search_for_cluster(clustered_list, index[0])) is not None:
                clustered_list[idx].extend([index[0], index[1]])
            elif (idx := search_for_cluster(clustered_list, index[1])) is not None:
                clustered_list[idx].extend([index[0], index[1]])
            else:
                clustered_list.append([index[0], index[1]])

    if len(clustered_list) > 0:
        idx_group = []
        new_targets = []
        for cluster in clustered_list:
            box_group = []
            for idx in set(cluster):
                box_group.append(results[idx])
                idx_group.append(idx)

            grouped = group_boxes(box_group)
            new_targets.append(grouped)
        new_targets = np.array(new_targets)
        results = np.delete(results, idx_group, axis=0)
        results = np.concatenate([new_targets, results], axis=0)
    return results


def search_for_cluster(clustered, idx):
    for i, cluster in enumerate(clustered):
        if idx in cluster:
            return i
    return


def group_boxes(boxes):
    boxes = np.array(boxes)
    conf = boxes[:, 1].max()
    x0 = boxes[:, 3].min()
    y0 = boxes[:, 4].min()
    x1 = boxes[:, 5].max()
    y1 = boxes[:, 6].max()
    return [0, conf, 0, x0, y0, x1, y1]
